/**
 * @file memory.c
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All Rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */

#include <stdlib.h>
#include "acl/acl_base.h"
#include "acl/acl_rt.h"
#include "log_inner.h"
#include "runtime/dev.h"
#include "runtime/mem.h"

#define POLICY_MASK_BIT 0xFF
#ifdef LITEOS_PRINT
#define APP 33
#endif
struct PageType {
  uint32_t aclPageType;
  uint32_t rtPageTypeBit;
};

static struct PageType g_pageMap[] = {
    {ACL_MEM_MALLOC_HUGE_FIRST, RT_MEMORY_POLICY_HUGE_PAGE_FIRST},
    {ACL_MEM_MALLOC_HUGE_ONLY, RT_MEMORY_POLICY_HUGE_PAGE_ONLY},
    {ACL_MEM_MALLOC_NORMAL_ONLY, RT_MEMORY_POLICY_DEFAULT_PAGE_ONLY},
    {ACL_MEM_MALLOC_HUGE_FIRST_P2P, RT_MEMORY_POLICY_HUGE_PAGE_FIRST_P2P},
    {ACL_MEM_MALLOC_HUGE_ONLY_P2P, RT_MEMORY_POLICY_HUGE_PAGE_ONLY_P2P},
    {ACL_MEM_MALLOC_NORMAL_ONLY_P2P, RT_MEMORY_POLICY_DEFAULT_PAGE_ONLY_P2P}};

aclError GetAlignedSize(const size_t size, size_t *alignedSize) {
  const size_t DATA_MEMORY_ALIGN_SIZE = 32UL;
  if ((size + (DATA_MEMORY_ALIGN_SIZE * 2UL)) < size) {
    return ACL_ERROR_INVALID_PARAM;
  }
  *alignedSize = (size + (DATA_MEMORY_ALIGN_SIZE * 2UL) - 1UL) / DATA_MEMORY_ALIGN_SIZE * DATA_MEMORY_ALIGN_SIZE;
  return ACL_SUCCESS;
}

static int Compare(const void *a, const void *b) {
  return (int)((*(struct PageType*)a).aclPageType - (*(struct PageType*)b).aclPageType);
}

aclError aclrtMalloc(void **devPtr, size_t size, aclrtMemMallocPolicy policy) {
  size_t alignedSize = 0;
  if ((devPtr == NULL) || (size == 0UL) || (GetAlignedSize(size, &alignedSize) != ACL_SUCCESS)) {
    ACL_LOG_ERROR("%s", devPtr == NULL? "devPtr is null" : "size invalid");
    return ACL_ERROR_INVALID_PARAM;
  }

  uint32_t flags = RT_MEMORY_DEFAULT;
  uint32_t memType = ACL_MEM_TYPE_LOW_BAND_WIDTH | ACL_MEM_TYPE_HIGH_BAND_WIDTH;
  if (((uint32_t)policy & memType) == memType) {
    ACL_LOG_WARN("policy can't contain both low and high width type.");
  } else if (((uint32_t)policy & ACL_MEM_TYPE_LOW_BAND_WIDTH) == ACL_MEM_TYPE_LOW_BAND_WIDTH) {
    flags = RT_MEMORY_DDR;
  } else if (((uint32_t)policy & ACL_MEM_TYPE_HIGH_BAND_WIDTH) == ACL_MEM_TYPE_HIGH_BAND_WIDTH) {
    flags = RT_MEMORY_HBM;
  }

  static size_t pageMapLen = sizeof(g_pageMap) / sizeof(g_pageMap[0]);
  static size_t pagePairLen = sizeof(g_pageMap[0]);
  struct PageType pagePair = {(uint32_t)policy & POLICY_MASK_BIT, 0};
  struct PageType *searchRet = (struct PageType*)bsearch(&pagePair, g_pageMap, pageMapLen, pagePairLen, Compare);
  if (searchRet == NULL) {
    ACL_LOG_ERROR("invalid policy!");
    return ACL_ERROR_INVALID_PARAM;
  }
  flags |= (*searchRet).rtPageTypeBit;
  return rtMalloc(devPtr, alignedSize, flags, (uint16_t)(APP));
}

aclError aclrtFree(void *devPtr) {
  if (devPtr == NULL) {
      ACL_LOG_ERROR("devPtr is null.");
      return ACL_ERROR_INVALID_PARAM;
  }
  return rtFree(devPtr);
}

aclError MemcpyKindTranslate(const aclrtMemcpyKind kind, rtMemcpyKind_t *rtKind) {
  switch (kind) {
    case ACL_MEMCPY_HOST_TO_DEVICE: {
      *rtKind = RT_MEMCPY_HOST_TO_DEVICE;
      break;
    }
    case ACL_MEMCPY_DEVICE_TO_DEVICE: {
      *rtKind = RT_MEMCPY_DEVICE_TO_DEVICE;
      break;
    }
    case ACL_MEMCPY_DEVICE_TO_HOST: {
      *rtKind = RT_MEMCPY_DEVICE_TO_HOST;
      break;
    }
    case ACL_MEMCPY_HOST_TO_HOST: {
      *rtKind = RT_MEMCPY_HOST_TO_HOST;
      break;
    }
    default: {
      return ACL_ERROR_INVALID_PARAM;
    }
  }
  return ACL_SUCCESS;
}

aclError aclrtMemcpy(void *dst, size_t destMax, const void *src, size_t count, aclrtMemcpyKind kind) {
  if (dst == NULL || src == NULL) {
    ACL_LOG_ERROR("%s",  dst == NULL? "dst is null": "src is null.");
    return ACL_ERROR_INVALID_PARAM;
  }
  rtMemcpyKind_t rtKind = RT_MEMCPY_RESERVED;
  const aclError ret = MemcpyKindTranslate(kind, &rtKind);
  if (ret != ACL_SUCCESS) {
    ACL_LOG_INNER_ERROR("invalid kind[%d]", (int32_t)(kind));
    return ret;
  }
  return rtMemcpy(dst, destMax, src, count, rtKind);
}

aclError aclrtMemset(void *devPtr, size_t maxCount, int32_t value, size_t count) {
  if (devPtr == NULL) {
      ACL_LOG_ERROR("devPtr is null.");
      return ACL_ERROR_INVALID_PARAM;
  }
  return rtMemset(devPtr, maxCount, (uint32_t)(value), count);
}

aclError aclrtGetMemInfo(aclrtMemAttr attr, size_t *free, size_t *total) {
  if (free == NULL || total == NULL) {
    ACL_LOG_ERROR("%s", free == NULL? "free is null" : "total is null.");
    return ACL_ERROR_INVALID_PARAM;
  }
  return rtMemGetInfoEx((rtMemInfoType_t)(attr), free, total);
}